import numpy as np
import torch
from mmdet3d.utils import vis_bev_pc
from mmdet3d.core import LiDARInstance3DBoxes

if __name__ == '__main__':


    a = np.array([[ 7.21755947e+03, -1.58339532e+03,  2.12940623e+02,
                2.19935703e+00,  4.80318069e+00,  1.81651068e+00,
                3.06824152e+00],
        [ 7.19642166e+03, -1.56225650e+03,  2.12097350e+02,
                1.91852856e+00,  4.45630503e+00,  1.53911805e+00,
                -1.79231252e-01],
        [ 7.19987757e+03, -1.59062623e+03,  2.12072546e+02,
                2.02242851e+00,  4.69164848e+00,  1.54427946e+00,
                -9.18192033e-02],
        [ 7.19360422e+03, -1.53284993e+03,  2.12245252e+02,
                2.08939242e+00,  4.71750259e+00,  1.55200136e+00,
                -2.68660362e-01],
        [ 7.23920300e+03, -1.54189480e+03,  2.14992504e+02,
                3.02812767e+00,  9.60439301e+00,  3.75812030e+00,
                -7.47424805e-02],
        [ 7.23048307e+03, -1.61195377e+03,  2.13452069e+02,
                2.92561388e+00,  7.21870804e+00,  2.50442243e+00,
                -3.13100008e+00],
        [ 7.23860939e+03, -1.55590835e+03,  2.14165083e+02,
                3.19131970e+00,  9.58074760e+00,  3.28009558e+00,
                -1.88764422e-02],
        [ 7.18255761e+03, -1.50126192e+03,  2.12236429e+02,
                2.02867913e+00,  4.46814919e+00,  1.64741468e+00,
                -3.97461186e-01],
        [ 7.20250681e+03, -1.62780220e+03,  2.12078538e+02,
                1.95667470e+00,  4.46179199e+00,  1.55260909e+00,
                -9.72686459e-02]])
    b = np.array([[ 7.19990129e+03, -1.59055269e+03,  2.12081712e+02,
         2.00786470e+00,  4.64698197e+00,  1.53390438e+00,
        -1.04685076e-01],
       [ 7.21758908e+03, -1.58333923e+03,  2.12954851e+02,
         2.16209585e+00,  4.80730539e+00,  1.79475591e+00,
         3.06364603e+00],
       [ 7.19642329e+03, -1.56231634e+03,  2.12101749e+02,
         1.97239324e+00,  4.53742944e+00,  1.55876896e+00,
        -1.95349350e-01],
       [ 7.20244790e+03, -1.62767280e+03,  2.12031252e+02,
         1.95819659e+00,  4.45224008e+00,  1.49711425e+00,
        -8.99051924e-02],
       [ 7.19363984e+03, -1.53284061e+03,  2.12204652e+02,
         2.07945002e+00,  4.75960343e+00,  1.50752014e+00,
        -2.79273186e-01],
       [ 7.23053817e+03, -1.61184225e+03,  2.13457932e+02,
         3.00906245e+00,  7.19318044e+00,  2.51191433e+00,
        -3.13559504e+00],
       [ 7.23866558e+03, -1.55586725e+03,  2.14241326e+02,
         3.05204789e+00,  9.55380927e+00,  3.62349188e+00,
        -2.15315755e-02],
       [ 7.23924602e+03, -1.54183317e+03,  2.15028666e+02,
         3.11392172e+00,  9.78852850e+00,  3.86400658e+00,
        -8.19391052e-02],
       [ 7.18248913e+03, -1.50099450e+03,  2.12271558e+02,
         2.00684096e+00,  4.37490340e+00,  1.62549114e+00,
        -3.98821820e-01]])

#     a = np.array([[ 7.23860939e+03, -1.55590835e+03,  2.14165083e+02,  3.19131970e+00,
#             9.58074760e+00,  3.28009558e+00, -1.55191988e+00]])
#     b = np.array([[ 7.23866558e+03, -1.55586725e+03,  2.14241326e+02,  3.05204789e+00,
#             9.55380927e+00,  3.62349188e+00, -1.54926475e+00]])
    a = a[[6], :]
    b = b[[6], :]
    center = a[:,:3].copy()
    a[:,:3] = 0
    b[:,:3] -= center # BUG: IoU will be wrong if do not normalize center
    a = LiDARInstance3DBoxes(a, box_dim=7, with_yaw=True, origin=(0.5, 0.5, 0.5))
    b = LiDARInstance3DBoxes(b, box_dim=7, with_yaw=True, origin=(0.5, 0.5, 0.5))
    a.tensor = a.tensor.cuda()
    b.tensor = b.tensor.cuda()
    iou = LiDARInstance3DBoxes.overlaps_bev(a, b)
    print('IoU from MMDet3D: ', iou)
    
    pc = np.array([[-10,-10, 0], [10, 10, 0]], dtype=np.float32)
    vis_bev_pc(pc, a, b, 'iou_debug_with_yaw_correction.png', './figs', figsize=(10, 10))